import io

import cv2
import numpy as np

from src import utils


def exam_num_recognition(image):
    """
    通过考生号图片识别考生号
    :param image: 考生号图片
    :return exam_num: 识别的考生号
    """
    thresh_img = get_exam_num_area(image)
    line_y_height, first_line_y = get_exam_num_height(thresh_img)
    # 腐蚀与膨胀
    kernel = np.ones((4, 4), np.uint8)
    result_image = cv2.morphologyEx(thresh_img, cv2.MORPH_ERODE, kernel, iterations=1)
    result_image = cv2.GaussianBlur(result_image, (5, 5), 0)
    result_image = cv2.morphologyEx(result_image, cv2.MORPH_DILATE, kernel, iterations=3)
    result_image = cv2.morphologyEx(result_image, cv2.MORPH_CLOSE, kernel, iterations=3)
    _, result_image = cv2.threshold(result_image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    # utils.show_image(result_image)

    # 学生填充考号的识别结果
    exam_num = io.StringIO()
    contours, _ = cv2.findContours(result_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours, boundingBoxes = utils.sort_contours(contours, utils.LEFT_TO_RIGHT)
    # res = cv2.drawContours(image, contours, -1, (0, 0, 255), 2)
    # utils.show_image(res)
    boxes = []
    for box in boundingBoxes:
        (x, y, w, h) = box
        diff_w = utils.selected_box_w - w
        diff_h = utils.selected_box_h - h
        if diff_w > 0:
            x -= int(diff_w / 2)
            x = max(0, x)
            w = utils.selected_box_w
        if diff_h > 0:
            y -= int(diff_h / 2)
            y = max(0, y)
            h = utils.selected_box_h
        box = (x, y, w, h)
        boxes.append(box)

    selected = []
    for box in boxes:
        center_x, center_y = utils.get_center_of_box(box)
        flag = True
        for s in selected:
            if utils.is_point_in_box(center_x, center_y, s):
                flag = False
                break
        if flag:
            selected.append(box)
            (x, y, w, h) = box
            exam_num.write(chr(int((y + h - first_line_y) / line_y_height) + ord('0')))
    exam_num = exam_num.getvalue()
    return exam_num


def get_exam_num_area(image):
    """
    识别考生号的填涂部分
    :param image: 考生号原图
    :return exam_num_area: 识别出的填涂区域的图片
    """
    result_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(result_image, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    kernel = np.ones((9, 9), np.uint8)
    result_image = cv2.morphologyEx(thresh, cv2.MORPH_GRADIENT, kernel)
    result_image = result_image - thresh
    contours, _ = cv2.findContours(result_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)
    # res = cv2.drawContours(image, contours[0], -1, (0, 0, 255), 2)
    # utils.show_image(res)
    peri = cv2.arcLength(contours[0], True)
    approx = cv2.approxPolyDP(contours[0], 0.06 * peri, True)
    x, y, w, h = cv2.boundingRect(approx)
    x += 2
    w -= 4
    box = (x, y, w, h)
    exam_num_area = utils.get_image_by_box(thresh, box)
    return exam_num_area


def get_exam_num_height(image):
    """
    获取答题卡填涂区域的行间距
    :param image: 答题卡填涂区域的图片
    :return line_y_height: 填涂区域的行间距, first_line_y: 第一行的y坐标
    """
    # 膨胀
    kernel = np.ones((3, 3), np.uint8)
    result_image = cv2.morphologyEx(image, cv2.MORPH_DILATE, kernel, iterations=1)
    result_image = cv2.GaussianBlur(result_image, (5, 5), 3)
    # utils.show_image(result_image)

    # 第一行待填充考号的中心的x轴坐标
    first_line_center_y = None
    # 第二行待填充考号的中心的x轴坐标
    second_line_center_y = None
    # 第一行待填充考号的底部边缘的y坐标
    first_line_bottom_y = None
    first_line_y = None
    contours, _ = cv2.findContours(result_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # res = cv2.drawContours(image.copy(), contours, -1, (0, 0, 255), 2)
    # utils.show_image(res)
    contours, _ = utils.sort_contours(contours, utils.TOP_TO_BOTTOM)
    # res = cv2.drawContours(origin_image, contours, -1, (0, 0, 255), 2)
    # utils.show_image(res)
    for c in contours:
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.06 * peri, True)
        (x, y, w, h) = cv2.boundingRect(approx)
        center_y = (2 * y + h) / 2

        if h > 10:
            if first_line_center_y is None:
                # temp_image = utils.get_image_by_box(image, (x, y, w, h))
                # utils.show_image(temp_image)
                first_line_center_y = center_y
                first_line_bottom_y = y + h
                first_line_y = y

            if center_y > first_line_bottom_y and second_line_center_y is None:
                # temp_image = utils.get_image_by_box(image, (x, y, w, h))
                # utils.show_image(temp_image)
                second_line_center_y = center_y
                break

    # print(type(second_line_center_y))

    line_y_height = second_line_center_y - first_line_center_y
    return line_y_height, first_line_y


if __name__ == "__main__":
    image = cv2.imread("../pic/exam_num/info.jpg")
    exam_num = exam_num_recognition(image)
    print(exam_num)
